[ExecuTorch][WebGPU] Add select/boolean ops (where, scalar compares, logical_not) to the WebGPU backend#20925
[ExecuTorch][WebGPU] Add select/boolean ops (where, scalar compares, logical_not) to the WebGPU backend#20925JCNTH wants to merge 4 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20925
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
SS-JIA
left a comment
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
Stack from ghstack (oldest at bottom):
Add
aten.where.self, scalar comparisons (aten.{eq,ne,le,ge,lt,gt}.Scalar), andaten.logical_not.defaulthandlers for the training tail. The scalar compares andlogical_notare generated from ONE templated byte-packed-bool shader (runtime/ops/boolean_op/) rather than per-op hand-written kernels — mirroring thebinary_opcodegen and Vulkan's grouping of comparison ops as${OPERATOR}variants (backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml).Key changes:
runtime/ops/boolean_op/— oneboolean_op.wgsltemplate +boolean_op.yaml(variantscompare_{eq,ne,le,ge,lt,gt}+logical_not) expand to the per-variant*_wgsl.h;BooleanOp.cppholds one shareddispatch_bool_ophelper (pack 4 bools peru32word, one thread per word) plus the sevenWEBGPU_REGISTER_OPs.runtime/ops/where/— ternary select, kept as its own kernel/handler (non-bool-family, as in Vulkan).op_registry.py— register the net-newaten.{ne,lt,le,ge,gt}.Scalar+aten.logical_not.defaultOpFeatures(eq.Scalaralready on master).CMakeLists.txtWEBGPU_SRCS— wireboolean_op/BooleanOp.cpp.aten.where.selfalready had VulkanOpFeatures; the scalar-compare +logical_notregistrations are additive to the shared Vulkan registry.Co-authored-with: Claude Code.
@exported-using-ghexport
Differential Revision: D111755122
Differential Revision: D111755122